1829G - Hits Different - CodeForces Solution

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define vi vector<ll>
#define vvi vector<vi>
#define pii pair<ll,ll>
#define vii vector<pii> 
#define rep(i,a,b) for(ll i=a;i<b;i++)
#define rev(i,a,b) for(ll i=a;i>=b;i--)
#define ff first 
#define ss second 
#define all(v) v.begin(), v.end()
#define srt(a) sort(a,a+n);
#define pb push_back
#define setBits(x) __builtin_popcountll(x)
#define yes cout << "YES" << endl;
#define no cout << "NO" << endl;
# define in(n) ll n; cin>>n;
# define str(s) string s; cin>>s;
# define ain(a,n) ll a[n]; rep(i,0,n) cin>>a[i];
# define aout(a,n) rep(i,0,n) cout<<a[i]<<" "; cout<<endl;
#define Reset(a, b) memset(a, b, sizeof(a));
# define en cout<<endl;
#define max_ele(a) *max_element(all(a));
#define min_ele(a) *min_element(all(a));

const int N = 1e5+5, MOD = 1e9+7;
# define LENGTH 100005

ll getGCD(ll a,ll b){ return b == 0 ? a : getGCD(b,a%b);}
ll getLCM(ll a,ll b){ return (a/getGCD(a,b))*b;}

ll powerM(ll x, ll  y){  
    if(y<0ll)
    return 0;
    ll result = 1ll;     
    x = x % MOD; 
    if (x == 0ll) return 0; 
    while (y > 0ll) {  
        if (y & 1ll)  
        result = (result*x) % MOD;  
        y = y>>1ll; 
        x = (x*x) % MOD;  
    }  
    return result;  
}

ll addM(ll x, ll y){
    x%=MOD;
    y%=MOD;
    return (x+y)%MOD;
}
 
ll mulM(ll x, ll y){
    return ((x%MOD) * 1ll * (y%MOD)) % MOD;
}

ll modeExp(ll x, ll y){
    if(y==0)
    return 1ll;
    ll ans = 1ll;
    x%=MOD;
    while(y){
        if(y & 1) ans = mulM(ans, x);
        x = mulM(x, x);
        y >>= 1ll;
    }
    return ans;
}
 
ll inviM(ll x){
    return modeExp(x, MOD - 2);
}
 
ll divideM(ll x, ll y){
    return mulM(x, inviM(y));
}

ll fact[LENGTH];
void fillfct (){
    fact[0] = 1;
    fact[1] = 1;
    for (int i = 2; i < LENGTH; i++)
    fact[i] = (fact[i - 1] * i)%MOD;
    return;
}

ll nCr(ll n, ll k){
    if(n<k) return 0;
    if(k<0) return 0;
    if(n==0) return 1;
    return divideM(fact[n], mulM(fact[k], fact[n - k]));
    //return (fact[n]/(fact[n-k]))/fact[k];
}

ll nPr(ll n,ll r,ll m=MOD){
    return divideM(fact[n],fact[n-r]);
}

ll subM(ll A,ll B){
    return (A-B+MOD)%MOD;
}

bool sortbysecond(const pair<ll,ll>&a,const pair<ll,ll>&b){return (a.second < b.second);}

// for (int i=0; i<n; i++)     // sum of all subarrays
// result += (arr[i] * (i+1) * (n-i));
// sum+= [((i + 1) * (n – i) + 1) / 2]*arr[i]   == for all odd length subarrays

// vector<ll> printPrimeFactors(ll n) {
//    vector<ll> ans;
//    while (n%2 == 0){
//       ans.pb(2);
//       n = n/2;
//    }
//    for (int i = 3; i <= sqrtl(n); i = i+2){
//       while (n%i == 0){
//          ans.pb(i);
//          n = n/i;
//       }
//    }
//    if (n > 2)
//    ans.pb(n);
//    return ans;
// }

// ll countBits(ll n){return (ll)log2(n)+1;}

// int countDivisors(int n){
//     int cnt = 0;
//     for (int i = 1; i <= sqrt(n); i++) {
//         if (n % i == 0) {
//             if (n / i == i) cnt++;
//             else cnt = cnt + 2;
//         }
//     }
//     return cnt;
// }

// bool PowerOfTwoH(ll n){if(n==0) return false; return (ceil(log2(n)) == floor(log2(n)));}

// int primes[1000009];
// vector<int> pr;
// void sieve(int n){
//     for(int i=2;i<=N;i++){
//        if(primes[i]==0){
//         pr.push_back(i);
//         for(int j=i*i;j<N;j+=i){
//             primes[j]=1;
//         }
//        }
//        primes[i]^=1; 
//     }
// }

// int leastDiv[1000009];
// void sieveDivisors(int n){
//     memset(leastDiv,-1,sizeof(leastDiv));
//     leastDiv[1]=1;
//     for(int i=2;i<=n;i++){
//        if(leastDiv[i]==-1){
//         for(int j=i;j<=n;j+=i){
//             if(leastDiv[j]==-1) leastDiv[j]=i;
//         }
//        }
//     }
// }

// vector<int> bfs(int nodes,vector<int> adj[]){
//     int vis[nodes]={0};
//     vis[0]=1;
//     queue<int>q;
//     q.push(0);
//     vector<int> ans;
//     while(!q.empty()){
//         int node = q.front();
//         q.pop();
//         ans.push_back(node);

//         for(auto it : adj[node]){
//             if(!vis[it]){
//                 vis[it]=1;
//                 q.push(it);
//             }
//         }
//     }
//     return ans;
// }

// void f(int node,vector<int> adj[],int vis[],vector<int> &ans){
//     vis[node]=1;
//     ans.push_back(node);
//     for(auto it : adj[node]){
//         if(!vis[it]){
//             f(it,adj,vis,ans);
//         }
//     }
// }

// vector<int> DFS(int nodes,vector<int> adj[]){
//     int vis[nodes]={0};
//     int start=0;
//     vector<int>ans;
//     f(start,adj,vis,ans);
//     return ans;
// }

// vector<ll>p(n,a[0]);
// vector<ll>s(n,a[n-1]);
// rep(1,n)
// p[i]=p[i-1]+a[i];
// for(ll i=n-2;i>=0;i--)
// s[i]=s[i+1]+a[i];
// vector<int>temp=primeFactors(n); in int main()


int main(){
	ios_base::sync_with_stdio(0);
	cin.tie(0);
    // #ifndef ONLINE_JUDGE
    // freopen("input.txt", "r", stdin);
    // freopen("output.txt", "w", stdout);
    // #endif
	// fillfct();
    // sieve(1000005);
    vi dp(1e6+1); dp[0]=0;dp[1]=1;dp[2]=5;dp[3]=10; ll r=3,cnt=0;                
    rep(i,4,1e6+1){                                                           //    1
        dp[i] = i*i;                                                        //     2  3
        ll left_ele = (r*(r-1)/2)+1;                                      //      4  5  6
        ll upper_ele = (((r-1)*(r-2))/2)+1;                              //     7   8  9  10
        ll upper_to_upper_ele = (((r-2)*(r-3))/2)+1;                              
        ll upper_row=r-1;                                               
        ll d=i-left_ele;

        if(d==upper_row){ 
        dp[i] += dp[left_ele-1];                  // right most element
        } 

        else if(d==0){                                       
        dp[i] += dp[upper_ele];                   // left most element
        } 

        else{
            dp[i] += dp[upper_ele+d] + dp[upper_ele+d-1] - dp[upper_to_upper_ele+d-1];
        }

        cnt++;

        if(cnt==r){
            cnt=0;
            r++;
        }
    }
    
	ll copy_nhi_kiya_mene;
    cin >> copy_nhi_kiya_mene;
	while (copy_nhi_kiya_mene--){
    in(n) cout<<dp[n]<<endl;
    }
}